from fastapi import APIRouter, WebSocket

router = APIRouter()


class WebsocketHeader:
    ws_list = []

    def __init__(self, ws):
        self.ws_list.append(ws)

    async def send(self, ws, data):
        await ws.send_text(f"Message text was: {data}")

    async def broadcast(self, data):
        for ws in self.ws_list:
            await self.send(ws, data)

    async def accept(self, ws):
        await ws.accept()
        self.ws_list.append(ws)


ws_header = WebsocketHeader(ws=None)


@router.websocket("/kline")
async def websocket_endpoint(websocket: WebSocket):
    # 等待 web 连接
    await websocket.accept()
    while True:
        data = await websocket.receive_text()
        await websocket.send_text(f"Message text was: {data}")
    # websockets.


@router.websocket("/kline2")
async def websocket2_endpoint(ws: WebSocket):
    # 等待 web 连接
    await ws_header.accept(ws)
    while True:
        data = await ws.receive_text()
        # 广播数据
        await ws_header.broadcast(data)


# @router.web("/kline2/{room}")
# async def websocket2_endpoint(ws: WebSocket):
#     # 等待 web 连接
#     await ws_header.accept(ws)
#     while True:
#         data = await ws.receive_text()
#         # 广播数据
#         await ws_header.broadcast(data)


